'''
匈牙利算法
'''
# vim: expandtab:ts=4:sw=4
from __future__ import absolute_import
import numpy as np
# from sklearn.utils.linear_assignment_ import linear_assignment
from scipy.optimize import linear_sum_assignment as linear_assignment
from . import kalman_filter
import configuration as cfg

INFTY_COST = 1e+5


def min_cost_matching(
        distance_metric, tracks, detections, track_indices=None,
        detection_indices=None):
    """
    解决线性分配问题，

    Parameters
    ----------
    distance_metric :
        一个回调函数(tracks, dets, track_indices, detection_indices)
        该回调函数参数(tracks列表，detections列表，N个轨迹索引列表，M个检测索引列表)
        该函数返回 NxM 维代价矩阵，其中元素 (i, j)
        是给定轨道索引中的第 i 个轨道与给定 detection_indices 中的
        第 j 个检测之间的关联成本。

    tracks : List[track.Track]
        A list of predicted tracks at the current time step.
    detections : List[detection.Detection]
        A list of detections at the current time step.
    track_indices : List[int]
        List of track indices that maps rows in `cost_matrix` to tracks in
        `tracks` (see description above).
    detection_indices : List[int]
        List of detection indices that maps columns in `cost_matrix` to
        detections in `detections` (see description above).
detections
    Returns
    -------
    (List[(int, int)], List[int], List[int])
        Returns a tuple with the following three entries:
        * A list of matched track and detection indices.
        * A list of unmatched track indices.
        * A list of unmatched detection indices.

    """
    #=========先检测是不是为空================
    if track_indices is None:
        track_indices = np.arange(len(tracks))
    if detection_indices is None:
        detection_indices = np.arange(len(detections))
    #==========为空则返回空===================
    if len(detection_indices) == 0 or len(track_indices) == 0:
        return [], track_indices, detection_indices
    #=================计算代价矩阵=============
    cost_matrix = distance_metric(
        tracks, detections, track_indices, detection_indices)
    cost_matrix[cost_matrix>cfg.matching_threshold] = cfg.matching_threshold+1e-5
    #执行匈牙利算法，得到匹配成功的索引对，行索引为tracks的索引，列索引为
    #detections的索引
    row_indices, col_indices = linear_assignment(cost_matrix)
    matches, unmatched_tracks, unmatched_detections = [], [], []
    #找出未matched的detctions
    for col, detection_idx in enumerate(detection_indices):
        if col not in col_indices:
            unmatched_detections.append(detection_idx)
    #找出未匹配的tracks
    for row, track_idx in enumerate(track_indices):
        if row not in row_indices:
            unmatched_tracks.append(track_idx)
    #遍历匹配的(track,detection)索引对
    for row, col in zip(row_indices, col_indices):
        track_idx = track_indices[row]
        detection_idx = detection_indices[col]
        #如果相应的cost>大于max_distance，也视为unmatched
        if cost_matrix[row, col] > cfg.matching_threshold:
            unmatched_tracks.append(track_idx)
            unmatched_detections.append(detection_idx)
        else:
            matches.append((track_idx, detection_idx))
    return matches, unmatched_tracks, unmatched_detections


def matching_cascade(
        distance_metric, tracks, detections,
        track_indices=None, detection_indices=None):
    """
    匹配tracks和detections。
    挑选出所有的confirmed tracks，优先让age较小的tracks和detections匹配。
    然后才对age较大的tracks进行匹配。这使得在相同的外观特征和马氏距离的情况下，
    age较小的跟踪器更容易匹配上。
    Parameters
    ----------
    distance_metric : 一个回调函数，Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
        distance metric给出一个tracks和detections的列表，以及N个tracks的索引和M个检测索引的列表
        metric应该返回NxM维的代价矩阵，里面的元素(i,j)是给定tracks的索引中第i个track和给定检测的索引
        中的第j个detection之间的关联代价矩阵。
    tracks : List[track.Track]
        使用卡尔曼滤波，根据上一时刻预测到的当前时刻的tracks
    detections : List[detection.Detection]
        目标检测器检测到的当前时刻的detections
    track_indices : Optional[List[int]]
        将``cost_matrix''映射到``tracks''轨迹的索引列表（请参见上面的描述）。
        默认为所有tracks。
    detection_indices : Optional[List[int]]
        将 `cost_matrix` 中的列映射到 `detections` 中的检测的检测索引列表（参见上面的描述）。
        默认为所有detections。

    Returns
    -------
    (List[(int, int)], List[int], List[int])
        返回一个包含以下3个元素的元组:
        * 一个包含了匹配到的tracks和detections索引列表
        * 一个包含了未匹配到的track索引列表
        * A list of unmatched detection indices.
    """
    if track_indices is None:
        track_indices = list(range(len(tracks)))
    if detection_indices is None:
        detection_indices = list(range(len(detections)))
    # 由于还未开始匹配detections，所以所有的detection都是未匹配的
    unmatched_dets_indices = detection_indices
    matches = []
    #===========遍历不同的age，从age较小的开始匹配========================
    for level in range(cfg.max_age):
        if len(unmatched_dets_indices) == 0:  #检测列表为空
            break
        #=====挑选出对应age的跟踪器tracks，按照age排序，age小的在前面========
        track_indices_l = []#当前level的所有tracks索引
        for k in track_indices:
            if tracks[k].time_since_update == 1 + level:
                track_indices_l.append(k)

        if len(track_indices_l) == 0:  #如果当前level没有track，继续
            continue
        #========将跟踪器tracks和unmatched的detections进行匹配========
        #匈牙利算法用来将前一帧中的跟踪框tracks和当前帧的检测框detections进行关联，通过
        #外观信息(appearance information)和马氏距离(Mahalanobis distance)或者IOU
        #来计算代价矩阵cos matrix
        matches_indices, unmatched_tracks_indices, unmatched_dets_indices = \
            min_cost_matching(
                distance_metric=distance_metric,
                tracks=tracks,
                detections=detections,
                track_indices=track_indices_l,
                detection_indices=unmatched_dets_indices)
        matches += matches_indices
    #=========挑选出unconfirmed的tracks=============
    unmatched_tracks_indices = list(set(track_indices) - set(k for k, _ in matches))
    return matches, unmatched_tracks_indices, unmatched_dets_indices

def gate_cost_matrix(
        kf, cost_matrix, tracks, detections, track_indices, detection_indices,
        gated_cost=INFTY_COST, only_position=False):
    """
    根据卡尔曼滤波得到的状态分布，使代价矩阵中的不可行的项失效
    （令其为一个特别大的值）。这个值为gated_cost
    Parameters
    ----------
    kf : The Kalman filter.
    cost_matrix : ndarray
        NxM的cost matrix，N是tracks indices的数量，M是detection indeces的数量，所以，
        元素(i, j)是tracks[track_indices[i]]和detections[detection_indices]]之间的
        关联成本。
    tracks : List[track.Track]，当前帧的预测tracks的列表
    detections : List[detection.Detection]，当前帧detections列表
    track_indices : List[int]
    detection_indices : List[int]
    gated_cost : Optional[float]
        使代价矩阵中的不可行的项失效（令其为一个特别大的值）
    only_position : Optional[bool]
        如果为True,只考虑状态分布的位置xy
    Returns
    -------
    ndarray
        返回被修改的代价矩阵。

    """
    gating_dim = 2 if only_position else 4
    #==========马式距离的门控阈值===========
    gating_threshold = kalman_filter.chi2inv95[gating_dim]
    #=======detections的坐标格式转换，从(tl_x,tl_y,b_w,b_h)->(c_x,c_y,w/h,h)
    measurements = np.asarray(
        [detections[i].to_xyah() for i in detection_indices])
    for row, track_idx in enumerate(track_indices):
        track = tracks[track_idx]   #根据tracks索引读取track
        #获得measurements测量值/检测值到track的均值矩阵和协方差矩阵的马氏距离
        gating_distance = kf.gating_distance(
            track.mean, track.covariance, measurements, only_position)
        cost_matrix[row, gating_distance > gating_threshold] = gated_cost
    return cost_matrix
